#include "EncodeAudio.h"
/*

static int select_best_sample_rate(const AVCodec* codec)
{
    const int* p;
    int best_samplerate = 0;
    if(!codec->supported_framerates){
        return 44100;
    }
    p = codec->supported_samplerates;
    while(*p)
    {
        if(!best_samplerate || abs(44100 - *p) < abs(44100 - best_samplerate))
        {
            best_samplerate = *p;
        }
        p++;
    }
    return best_samplerate;
}

static int check_sample_fmt(const AVCodec* codec, enum AVSampleFormat sample_fmt)
{
    const enum AVSampleFormat* p  = codec->sample_fmts;

    while (*p != AV_SAMPLE_FMT_NONE)
    {
        if(*p == sample_fmt)
        {
            return 1;
        }
        p++;
    }
    return 0;
}

// static 私有
static int encode(AVCodecContext* ctx,AVFrame* frame,AVPacket* pkt,FILE* out)
{
    int ret = -1;
    ret = avcodec_send_frame(ctx,frame);
    if(ret < 0)
    {
        av_log(NULL,AV_LOG_ERROR,"failed to send frame to encoder");
        goto _END;
    }

    while(ret >= 0)
    {
        ret = avcodec_receive_packet(ctx,pkt);
        if(ret == AVERROR(EAGAIN) || ret == AVERROR_EOF)
        {
            return 0;
        }
        else if(ret < 0)
        {
            return -1;
        }
        fwrite(pkt->data,1,pkt->size,out);
        av_packet_unref(pkt);
    }
_END:
    return 0;
}

int encodeAudio(const char* dst)
{
    const AVCodec* codec = NULL;
    AVCodecContext* ctx = NULL;
    int ret = -1;
    FILE * f = NULL;
    AVPacket* pkt = NULL;
    AVFrame* frame = NULL;

    uint16_t* samples = NULL;
    //uint32_t* samples = NULL;

    av_log_set_level(AV_LOG_DEBUG);
    // 1.参数检查

    // 2.查找编码器(index、名字)
    codec = avcodec_find_encoder_by_name("libfdk_aac"); // 第三方的AAC
    //codec = avcodec_find_encoder(AV_CODEC_ID_AAC);

    if(!codec)
    {
        av_log(NULL,AV_LOG_ERROR,"don't find codec\n");
        goto _ERROR;
    }

    // 3.创建编码器上下文
    ctx = avcodec_alloc_context3(codec);
    if(!ctx)
    {
        av_log(NULL,AV_LOG_ERROR,"no memrory\n");
        goto _ERROR;
    }

    // 4.设置编码器参数
    ctx->bit_rate = 64000;
    ctx->sample_fmt = AV_SAMPLE_FMT_S16;    // 采样大小
    //ctx->sample_fmt = AV_SAMPLE_FMT_FLTP;    // 采样大小
    if(!check_sample_fmt(codec,ctx->sample_fmt))
    {
        av_log(NULL,AV_LOG_ERROR,"Encoder does not support sample format\n");
        goto _ERROR;
    }

    ctx->sample_rate = select_best_sample_rate(codec);
    av_channel_layout_copy(&ctx->ch_layout,&(AVChannelLayout)AV_CHANNEL_LAYOUT_STEREO);
    //av_channel_layout_copy(&ctx->ch_layout,&(AVChannelLayout)AV_CHANNEL_LAYOUT_MONO);

    // 5.编码器和编码器上下文绑定
    ret = avcodec_open2(ctx,codec,NULL);
    if(ret < 0)
    {
        av_log(ctx,AV_LOG_ERROR,"don't open codec : %s \n",av_err2str(ret));
        goto _ERROR;
    }

    // 6.创建输出文件
    f = fopen(dst,"wb");
    if(!f)
    {
        av_log(NULL,AV_LOG_ERROR,"don't open file: %s\n",dst);
        goto _ERROR;
    }

    // 7.创建AVFrame
    frame = av_frame_alloc();
    if(!frame)
    {
        av_log(NULL,AV_LOG_ERROR,"no memory\n");
        goto _ERROR;
    }

    frame->nb_samples = ctx->frame_size;
    frame->format = AV_SAMPLE_FMT_S16;
    //frame->format = AV_SAMPLE_FMT_FLTP;

    av_channel_layout_copy(&frame->ch_layout,&(AVChannelLayout)AV_CHANNEL_LAYOUT_STEREO);
    //av_channel_layout_copy(&frame->ch_layout,&(AVChannelLayout)AV_CHANNEL_LAYOUT_MONO);
    frame->sample_rate = ctx->sample_rate;

    ret = av_frame_get_buffer(frame,0);
    if(ret < 0)
    {
        av_log(NULL,AV_LOG_ERROR,"Could not allocate the video frame\n");
        goto _ERROR;
    }

    // 8.创建
    pkt = av_packet_alloc();
    if(!pkt)
    {
        av_log(NULL,AV_LOG_ERROR,"no memory\n");
        goto _ERROR;
    }

    // 9.生成音频内容（可通过抓取屏幕、摄像头录制）
    float t = 0;
    float tincr = 2*M_PI*440 / ctx->sample_rate;
    for(int i = 0;i < 200;i++)
    {
        ret = av_frame_make_writable(frame);

        if(ret < 0)
        {
            av_log(NULL,AV_LOG_ERROR,"Could not allocate space\n");
            goto _ERROR;
        }

        samples = (uint16_t*)frame->data[0];
        //samples = (uint32_t*)frame->data[0];

        for(int j = 0;j < ctx->frame_size;j++)
        {
            samples[2* j] = (int)(sin(t) * 10000);
            //samples[4* j] = (int)(sin(t) * 10000); // 数值sin(t) * 10000可能要该
            for(int k = 1;k < ctx->ch_layout.nb_channels;k++)
            {
                samples[2 * j + k] = samples[2 * j];
                //samples[4 * j + k] = samples[4 * j];
            }
            t+= tincr;
        }
        encode(ctx,frame,pkt,f);
    }

    // 10.编码,防止缓冲区中的数据残留
    encode(ctx,NULL,pkt,f);

_ERROR:
    if(ctx)
    {
        avcodec_free_context(&ctx);
    }

    if(frame)
    {
        av_frame_free(&frame);
    }

    if(pkt)
    {
        av_packet_free(&pkt);
    }

    if(f)
    {
        fclose(f);
    }

    return 0;
}

*/

#include "EncodeAudio.h"
#include <math.h>

// 添加ADTS头函数
static void add_adts_header(uint8_t* packet, int packet_size, int profile, int sample_rate, int channels) {
    int sample_rate_index = 0;

    // 采样率索引表
    switch (sample_rate) {
    case 96000: sample_rate_index = 0; break;
    case 88200: sample_rate_index = 1; break;
    case 64000: sample_rate_index = 2; break;
    case 48000: sample_rate_index = 3; break;
    case 44100: sample_rate_index = 4; break;
    case 32000: sample_rate_index = 5; break;
    case 24000: sample_rate_index = 6; break;
    case 22050: sample_rate_index = 7; break;
    case 16000: sample_rate_index = 8; break;
    case 12000: sample_rate_index = 9; break;
    case 11025: sample_rate_index = 10; break;
    case 8000: sample_rate_index = 11; break;
    case 7350: sample_rate_index = 12; break;
    default: sample_rate_index = 4; break; // 默认44100
    }

    int channel_config = channels;
    int frame_length = packet_size + 7; // ADTS头7字节

    // 构建ADTS头
    packet[0] = 0xFF; // syncword 8 bits
    packet[1] = 0xF1; // syncword 4 bits + ID 1 bit + layer 2 bits + protection_absent 1 bit
    packet[2] = ((profile - 1) << 6) | (sample_rate_index << 2) | (channel_config >> 2);
    packet[3] = ((channel_config & 3) << 6) | (frame_length >> 11);
    packet[4] = (frame_length >> 3) & 0xFF;
    packet[5] = ((frame_length & 7) << 5) | 0x1F;
    packet[6] = 0xFC; // buffer fullness
}

static int select_best_sample_rate(const AVCodec* codec)
{
    const int* p;
    int best_samplerate = 0;
    if(!codec->supported_samplerates){
        return 44100;
    }
    p = codec->supported_samplerates;
    while(*p)
    {
        if(!best_samplerate || abs(44100 - *p) < abs(44100 - best_samplerate))
        {
            best_samplerate = *p;
        }
        p++;
    }
    return best_samplerate;
}

static int check_sample_fmt(const AVCodec* codec, enum AVSampleFormat sample_fmt)
{
    const enum AVSampleFormat* p  = codec->sample_fmts;

    while (*p != AV_SAMPLE_FMT_NONE)
    {
        if(*p == sample_fmt)
        {
            return 1;
        }
        p++;
    }
    return 0;
}

// 选择最佳的采样格式
static enum AVSampleFormat select_best_sample_format(const AVCodec* codec)
{
    const enum AVSampleFormat* p = codec->sample_fmts;
    enum AVSampleFormat best_fmt = AV_SAMPLE_FMT_NONE;

    if (!p || *p == AV_SAMPLE_FMT_NONE) {
        return AV_SAMPLE_FMT_FLTP; // AAC编码器通常使用FLTP
    }

    // 优先选择FLTP（AAC常用）
    while (*p != AV_SAMPLE_FMT_NONE) {
        if (*p == AV_SAMPLE_FMT_FLTP) {
            return AV_SAMPLE_FMT_FLTP;
        }
        p++;
    }

    // 如果没有FLTP，返回第一个支持的格式
    return codec->sample_fmts[0];
}

// static 私有
static int encode(AVCodecContext* ctx, AVFrame* frame, AVPacket* pkt, FILE* out)
{
    int ret = -1;
    ret = avcodec_send_frame(ctx, frame);
    if(ret < 0)
    {
        av_log(NULL, AV_LOG_ERROR, "Failed to send frame to encoder: %s\n", av_err2str(ret));
        goto _END;
    }

    while(ret >= 0)
    {
        ret = avcodec_receive_packet(ctx, pkt);
        if(ret == AVERROR(EAGAIN) || ret == AVERROR_EOF)
        {
            return 0;
        }
        else if(ret < 0)
        {
            av_log(NULL, AV_LOG_ERROR, "Failed to receive packet from encoder: %s\n", av_err2str(ret));
            return -1;
        }

        // 添加ADTS头
        uint8_t adts_header[7];
        int profile = 1; // AAC LC profile
        if (ctx->sample_fmt == AV_SAMPLE_FMT_FLTP) {
            profile = 1; // AAC LC
        } else if (ctx->sample_fmt == AV_SAMPLE_FMT_S16) {
            profile = 1; // AAC LC
        }

        add_adts_header(adts_header, pkt->size, profile, ctx->sample_rate, ctx->ch_layout.nb_channels);

        // 先写入ADTS头
        fwrite(adts_header, 1, 7, out);
        // 再写入AAC数据
        fwrite(pkt->data, 1, pkt->size, out);

        av_log(NULL, AV_LOG_DEBUG, "Encoded packet: size=%d (with ADTS: %d), pts=%lld\n",
               pkt->size, pkt->size + 7, pkt->pts);
        av_packet_unref(pkt);
    }
_END:
    return 0;
}

int encodeAudio(const char* dst)
{
    const AVCodec* codec = NULL;
    AVCodecContext* ctx = NULL;
    int ret = -1;
    FILE * f = NULL;
    AVPacket* pkt = NULL;
    AVFrame* frame = NULL;

    float** samples = NULL;

    av_log_set_level(AV_LOG_DEBUG);

    // 1.参数检查
    if (!dst) {
        av_log(NULL, AV_LOG_ERROR, "Destination file path is NULL\n");
        goto _ERROR;
    }

    // 2.查找FFmpeg内部的AAC编码器
    codec = avcodec_find_encoder(AV_CODEC_ID_AAC);
    if(!codec)
    {
        av_log(NULL, AV_LOG_ERROR, "Don't find AAC codec\n");

        // 尝试其他AAC编码器
        codec = avcodec_find_encoder_by_name("aac");
        if (!codec) {
            av_log(NULL, AV_LOG_ERROR, "Don't find any AAC encoder\n");
            goto _ERROR;
        }
    }

    av_log(NULL, AV_LOG_INFO, "Using AAC codec: %s\n", codec->name);

    // 3.创建编码器上下文
    ctx = avcodec_alloc_context3(codec);
    if(!ctx)
    {
        av_log(NULL, AV_LOG_ERROR, "No memory for codec context\n");
        goto _ERROR;
    }

    // 4.设置编码器参数
    ctx->bit_rate = 128000; // 128kbps，AAC常用比特率

    // 选择最佳的采样格式
    ctx->sample_fmt = select_best_sample_format(codec);
    av_log(NULL, AV_LOG_INFO, "Selected sample format: %s\n", av_get_sample_fmt_name(ctx->sample_fmt));

    ctx->sample_rate = select_best_sample_rate(codec);
    av_log(NULL, AV_LOG_INFO, "Selected sample rate: %d\n", ctx->sample_rate);

    // 设置声道布局
    ret = av_channel_layout_copy(&ctx->ch_layout, &(AVChannelLayout)AV_CHANNEL_LAYOUT_STEREO);
    if (ret < 0) {
        av_log(NULL, AV_LOG_ERROR, "Failed to set channel layout: %s\n", av_err2str(ret));
        goto _ERROR;
    }

    // 设置帧大小
    ctx->frame_size = 1024; // AAC常用帧大小

    // 5.编码器和编码器上下文绑定
    ret = avcodec_open2(ctx, codec, NULL);
    if(ret < 0)
    {
        av_log(ctx, AV_LOG_ERROR, "Don't open codec: %s\n", av_err2str(ret));
        goto _ERROR;
    }

    // 6.创建输出文件
    f = fopen(dst, "wb");
    if(!f)
    {
        av_log(NULL, AV_LOG_ERROR, "Don't open file: %s\n", dst);
        goto _ERROR;
    }

    // 7.创建AVFrame
    frame = av_frame_alloc();
    if(!frame)
    {
        av_log(NULL, AV_LOG_ERROR, "No memory for frame\n");
        goto _ERROR;
    }

    frame->nb_samples = ctx->frame_size;
    frame->format = ctx->sample_fmt;
    frame->sample_rate = ctx->sample_rate;

    // 复制声道布局
    ret = av_channel_layout_copy(&frame->ch_layout, &ctx->ch_layout);
    if (ret < 0) {
        av_log(NULL, AV_LOG_ERROR, "Failed to copy channel layout to frame: %s\n", av_err2str(ret));
        goto _ERROR;
    }

    ret = av_frame_get_buffer(frame, 0);
    if(ret < 0)
    {
        av_log(NULL, AV_LOG_ERROR, "Could not allocate the audio frame data\n");
        goto _ERROR;
    }

    // 8.创建AVPacket
    pkt = av_packet_alloc();
    if(!pkt)
    {
        av_log(NULL, AV_LOG_ERROR, "No memory for packet\n");
        goto _ERROR;
    }

    // 9.生成音频内容（440Hz正弦波）
    float t = 0;
    float tincr = 2 * M_PI * 440.0f / ctx->sample_rate; // 440Hz正弦波
    int num_frames = 200; // 生成200帧

    av_log(NULL, AV_LOG_INFO, "Start encoding %d frames...\n", num_frames);

    for(int i = 0; i < num_frames; i++)
    {
        ret = av_frame_make_writable(frame);
        if(ret < 0)
        {
            av_log(NULL, AV_LOG_ERROR, "Frame is not writable: %s\n", av_err2str(ret));
            goto _ERROR;
        }

        // 根据采样格式处理数据
        if (ctx->sample_fmt == AV_SAMPLE_FMT_FLTP) {
            // FLTP格式：每个声道一个平面，每个样本是float
            samples = (float**)frame->data;

            for (int ch = 0; ch < ctx->ch_layout.nb_channels; ch++) {
                float* channel_data = samples[ch];
                float current_t = t;

                for (int j = 0; j < frame->nb_samples; j++) {
                    channel_data[j] = sinf(current_t) * 0.3f; // 减小幅度避免削波
                    current_t += tincr;
                }
            }
            t += tincr * frame->nb_samples;

        } else if (ctx->sample_fmt == AV_SAMPLE_FMT_S16) {
            // S16格式：交错存储
            int16_t* sample_data = (int16_t*)frame->data[0];
            int sample_index = 0;

            for (int j = 0; j < frame->nb_samples; j++) {
                int16_t sample_value = (int16_t)(sinf(t) * 8000.0f); // 减小幅度

                for (int ch = 0; ch < ctx->ch_layout.nb_channels; ch++) {
                    sample_data[sample_index++] = sample_value;
                }
                t += tincr;
            }
        } else {
            av_log(NULL, AV_LOG_ERROR, "Unsupported sample format: %s\n",
                   av_get_sample_fmt_name(ctx->sample_fmt));
            goto _ERROR;
        }

        // 设置帧的PTS
        frame->pts = i * frame->nb_samples;

        // 编码这一帧
        ret = encode(ctx, frame, pkt, f);
        if (ret < 0) {
            av_log(NULL, AV_LOG_ERROR, "Failed to encode frame %d\n", i);
            goto _ERROR;
        }

        // 显示进度
        if (i % 50 == 0) {
            av_log(NULL, AV_LOG_INFO, "Encoded %d/%d frames\n", i, num_frames);
        }
    }

    // 10.刷新编码器，处理缓冲区中剩余的数据
    av_log(NULL, AV_LOG_INFO, "Flushing encoder...\n");
    encode(ctx, NULL, pkt, f);

    av_log(NULL, AV_LOG_INFO, "Audio encoding completed successfully: %s\n", dst);

_ERROR:
    if(ctx)
    {
        avcodec_free_context(&ctx);
    }

    if(frame)
    {
        av_frame_free(&frame);
    }

    if(pkt)
    {
        av_packet_free(&pkt);
    }

    if(f)
    {
        fclose(f);
    }

    return (ret < 0) ? -1 : 0;
}
